import copy
import time

import numpy as np
import pybullet as p
from simutils.pybullet import make_object_visual_center

import constants as c
from robot import ROBOT
from world import WORLD


class SIMULATION:
    
    
    def __init__(self, index=-1) -> None:
        
        self.__index = index
        self.__fitness = None
        self.__position = self.__get_random_position()
        self.__world = WORLD(self.__position)
        self.__newFitnessArray = []
        self.__fitnessArray = []
        
        
    def __get_random_position(self):
        
        if c.mode == "fixed":
            return np.array([10, 0, 0.5])
        else:
            position = (np.random.random(3) * 2 - 1) * 20
            position[2] = 0.5
            return position
    
    
    def __simulate(self, mode, robot, task):
        
        p.connect(mode)
        p.setGravity(0, 0, -9.8)
        
        robot.initialize_simulation()
        self.__world.load_world()
        
        for _ in range(c.epochs):
            
            p.stepSimulation()

            result = robot.run(robot.tasks[task])
            
            if result == 1:
                break
            
    
        position = robot.get_info("position")
        orientation = robot.get_info("orientation")

        p.disconnect()
        
        return position, orientation
    
    
    def __step(self, robot):
        
        p.stepSimulation()
        make_object_visual_center(robot.id_)
        time.sleep(1/500)

    
    def train(self, generations, task):
        
        self.robot = ROBOT(task, None)
        position, orientation = self.__simulate(p.DIRECT, self.robot, task)
        self.__fitness = self.robot.tasks[task].get_fitness(position, orientation)
        
        for i in range(generations):
            
            new_robot = copy.deepcopy(self.robot)
            new_robot.tasks[task].mutate()
            position, orientation = self.__simulate(p.DIRECT, new_robot, task)
            new_fitness = new_robot.tasks[task].get_fitness(position, orientation)
            
            self.__newFitnessArray.append(new_fitness)
            self.__fitnessArray.append(self.__fitness)
            
            if new_fitness > self.__fitness:
                self.__fitness = new_fitness
                self.robot = new_robot
                
            
            print(f"第{i}代子代，fitness={self.__fitness}")
            
        self.robot.save_networks(self.__index)
        self.__newFitnessArray = np.array(self.__newFitnessArray)
        self.__fitnessArray = np.array(self.__fitnessArray)
        np.save(f"logs/{task}-fitness.npy", self.__fitnessArray)
        np.save(f"logs/{task}-newFitness.npy", self.__newFitnessArray)


    def show(self, task, rotation_path, locomotion_path, jumping_path):
        
        robot = ROBOT(task, None)
        
        if robot.tasks["rotation"] is not None:
            robot.tasks["rotation"].load_network(rotation_path)
            
        if robot.tasks["locomotion"] is not None:
            robot.tasks["locomotion"].load_network(locomotion_path)
            
        if robot.tasks["jumping"] is not None:
            robot.tasks["jumping"].load_network(jumping_path)
        
        p.connect(p.GUI)
        p.setGravity(0, 0, -9.8)
        
        self.__world.load_world()
        robot.initialize_simulation()
        
        if task != "all":
            
            while True:
                
                try:
                    self.__step(robot)
                    print(robot.get_info("orientation"))
                    robot.run(robot.tasks[task])
                except KeyboardInterrupt:
                    break
        
        else:
            while True:
                
                self.__step(robot)
                angle = robot.get_info("orientation")[2]
                position = robot.get_info("position")
                if get_angle_diff(angle, position, self.__position) > 1 / 180 * np.pi:
                    print(get_angle_diff(angle, position, self.__position))
                    result = robot.run(robot.tasks["rotation"])
                
                else:
                    distance = (position[0] - self.__position[0]) ** 2 + (position[1] - self.__position[1]) ** 2
                    for _ in range(10 * int(distance)):
                        self.__step(robot)

        robot.save_motor_values() 
        p.disconnect()
        
    
def get_angle_diff(angle, position, target):
    
    vector_one = np.array([np.cos(angle), np.sin(angle)])
    vector_two = np.array([target[0] - position[0], target[1] - position[1]])
    cos_theta = np.dot(vector_one, vector_two) / (vector_two[0] ** 2 + vector_two[1] ** 2) ** 0.5
    return np.arccos(cos_theta)